import numpy as np
import math
#import matplotlib.pyplot as plt

import rospy
from mavros_msgs.msg import GlobalPositionTarget, State, RCIn, PositionTarget
from mavros_msgs.srv import CommandBool, CommandTOL, SetMode
from geometry_msgs.msg import PoseStamped, Twist, TwistStamped, Point
from sensor_msgs.msg import Imu, NavSatFix, Temperature
from std_msgs.msg import Float32, Float64, String, Float32MultiArray, Int8
from nav_msgs.msg import Odometry

mav_yaw = 0.0
mav_pos = np.array([0, 0, 0])
is_initialize_pos = False

def road():
    x_path = []
    y_path = []
    bound = []
    z = 2
    l = 4
    x_path.append(0)
    y_path.append(0)
    for k in range(1, 1000):
        y_path.append(0)
        x_path.append(0.05*k)
        bound.append(l)

    for k in range(1000, 2000):
        y_path.append(-10 + 10 * math.cos((k - 1000) / 1000 * math.pi))
        x_path.append(50 + 10 * math.sin((k - 1000) / 1000 * math.pi))
        bound.append(l)

    for k in range(2000, 3000):
        y_path.append(-28+8*math.cos((k-2000)/1000*math.pi/2))
        x_path.append(50-8*math.sin((k-2000)/1000*math.pi/2))
        bound.append(l)

    for k in range(3000, 4000):
        y_path.append(-28 - 12 * math.sin((k - 3000) / 1000 * math.pi / 2))
        x_path.append(30 - 12 * math.cos((k - 3000) / 1000 * math.pi / 2 + math.pi))
        bound.append(l)

    for k in range(4000, 5000):
        y_path.append(-40)
        x_path.append(30-0.02*(k-4000))
        bound.append(l)

    for k in range(5000, 6000):
        y_path.append(-30 - 10 * math.cos((k - 5000) / 1000 * math.pi / 3*2))
        x_path.append(10 - 10 * math.sin((k - 5000) / 1000 * math.pi / 3*2))
        bound.append(l)

    for k in range(6000, 7000):
        y_path.append(-25 + math.sqrt(3)*(k-6000)*0.005)
        x_path.append(10-5*math.sqrt(3) + (k-6000)*0.005)
        bound.append(l)

    for k in range(7000, 8000):
        y_path.append(-23+5*math.sqrt(3)+4*math.sin(-1*math.pi/6+(k-7000)/1000*math.pi*2/3))
        x_path.append(15-7*math.sqrt(3)+4*math.sin(math.pi/3+(k-7000)/1000*math.pi*2/3))
        bound.append(l)

    for k in range(8000, 9000):
        y_path.append(-19+5*math.sqrt(3))
        x_path.append(15-7*math.sqrt(3) - (15-7*math.sqrt(3))/1000*(k-8000))
        bound.append(l)

    r = (-19 + 5 * math.sqrt(3)) / 2
    for k in range(9000, 10000):
        y_path.append(r+r*math.cos((k-9000)/1000*math.pi))
        x_path.append(r*math.sin((k-9000)/1000*math.pi))
        bound.append(l)

    x_path = np.array(x_path)*0.4
    y_path = np.array(y_path)*0.7

    return x_path, y_path, z, bound


def get_point(x_path, y_path, cur_x, cur_y, last_point):
    min = 10000
    point = last_point
    if last_point < 9000:
        end_point = last_point+1000
    else: 
        end_point = 9999
    for i in range(last_point, end_point):
        dis = np.linalg.norm([x_path[i] - cur_x, y_path[i] - cur_y])
        if dis < min:
            min = dis
            point = i
    return point


def run_with_v_control(x_path, y_path, z, bound, max_error, v, kp, v_max):
    global is_initialize_pos, mav_pos, mav_yaw
    current = 1# 当前所处的位置对应路径点
    last_error = 0
    error = 0
    kp = kp
    kd = 5
    error_list = []
    error_list.append(error)

    # mavros
    rospy.init_node('mavros_control')
    local_pose_sub = rospy.Subscriber("mavros/local_position/pose", PoseStamped, local_pose_callback)
    local_vel_pub = rospy.Publisher('/mavros/setpoint_raw/local', PositionTarget, queue_size=10)
    rate = rospy.Rate(10)
    while not rospy.is_shutdown():

        # 获取无人机当前位置

        cur_x = mav_pos[0]
        cur_y = mav_pos[1]
        cur_z = mav_pos[2]
        print(cur_x)
        print(cur_y)
      
        # 得到无人机上次控制时的位置对应路径点与误差
        last = current
        last_error = error

        # 得到当前路径点
        current = get_point(x_path, y_path, cur_x, cur_y, current)
        print(current)
        if current == 0:
            current = 1

        # 切向
        dir = np.array([x_path[current] - x_path[current - 1], y_path[current] - y_path[current - 1]])
        
        # 切向单位向量
        dir = dir / np.linalg.norm(dir)

        error = np.linalg.norm([cur_x - x_path[current], cur_y - y_path[current]]) # 误差大小
        error_list.append(error)

        v_cmd_dir = dir * v[current]  # 切向速度控制量

        perp = np.array([x_path[current] - cur_x, y_path[current] - cur_y])  # 法向
        perp = perp / np.linalg.norm(perp)  # 单位向量

        v_cmd_perp = kp * error * perp + kd * (error - last_error) * perp  # 法向控制量
        # v_cmd_perp = 0.8 * v_init[current] * error * perp + kd * (error - last_error) * perp
        v_cmd = v_cmd_perp + v_cmd_dir # 合控制量
        v_cmd = sat(v_cmd, v_max)

        kp_vl = 5  # 迭代率
        temp = kp_vl * (1 / (1 + np.exp(-error + bound[current] * 3 / 4)) - 0.5)
        if last > 20:
            if v[last] - temp < 2:
                v[last - 10:current] = 2
            else:
                v[last - 10:current] = v[last] - temp
        else:
            v[last:current] = v[last] - temp

        # 执行控制

        command_vel = construct_veltarget_ENU(v_cmd[0], v_cmd[1], z-cur_z, 0.0)
        print("here to pubulish once")
        local_vel_pub.publish(command_vel)

        # 终点
        if current > len(x_path) - 3:
            break
        rate.sleep()
    # rospy.spin()
    l_error = max_error
    max_error = max(error_list)
    return v, max_error, kp
'''
    if max_error > l_error:
        kp = 1 * (max_error - l_error) + kp
    else:
        kp = kp

    return v, max_error, kp
'''

def sat(v, vm):
    value_v = np.linalg.norm(v)
    if value_v > vm:
        v = v/value_v*vm
    return v


def construct_veltarget_ENU(E=0, N=0, U=0, yaw_rate=0):
    target_raw_pose = PositionTarget()
    target_raw_pose.header.stamp = rospy.Time.now()

    '''
    uint8 FRAME_LOCAL_NED = 1
    uint8 FRAME_LOCAL_OFFSET_NED = 7
    uint8 FRAME_BODY_NED = 8
    uint8 FRAME_BODY_OFFSET_NED = 9
    '''
    # 直接就ENU了
    target_raw_pose.coordinate_frame = 1

    target_raw_pose.velocity.x = E
    target_raw_pose.velocity.y = N
    target_raw_pose.velocity.z = U

    target_raw_pose.type_mask = PositionTarget.IGNORE_PX + PositionTarget.IGNORE_PY + PositionTarget.IGNORE_PZ \
                                + PositionTarget.IGNORE_AFX + PositionTarget.IGNORE_AFY + PositionTarget.IGNORE_AFZ \
                                + PositionTarget.FORCE + PositionTarget.IGNORE_YAW

    target_raw_pose.yaw_rate = yaw_rate
    return target_raw_pose


def local_pose_callback(msg):
    global is_initialize_pos, mav_pos, mav_yaw
    is_initialize_pos = True
    mav_pos = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z])

    q0, q1, q2, q3 = msg.pose.orientation.w, msg.pose.orientation.x, msg.pose.orientation.y, msg.pose.orientation.z
    mav_yaw = math.atan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))


def RC_main(epoch):
    x_path, y_path, z, bound = road()
    v = 2*np.ones([10001,1])
    max_error = 1
    kp = 1
    v_max = 8
    for i in range(epoch):
        v, max_error, kp = run_with_v_control(x_path, y_path, z, bound, max_error, v, kp, v_max)
    return


if __name__ == "__main__":
    RC_main(10)

